from matplotlib.quiver import QuiverKey, Quiver
from matplotlib.offsetbox import DrawingArea, AnchoredOffsetbox
from matplotlib.transforms import Bbox
import numpy as np
from .length_units import LenType, Pt, as_unit_value
from typing import Literal
from matplotlib import text as mtext
from matplotlib import collections as mcollections
from matplotlib.pyplot import rcParams
from matplotlib.axes import Axes

LocType = Literal['upper right', 'upper left', 'lower right',
                    'lower left', 'right', 'center left', 'center right',
                    'lower center', 'upper center', 'center', 1, 2, 3, 4,
                    5, 6, 7, 8, 9, 10]


class QuiverLegend(AnchoredOffsetbox):
    halign = {'N': 'center', 'S': 'center', 'E': 'left', 'W': 'right'}
    valign = {'N': 'bottom', 'S': 'top', 'E': 'center', 'W': 'center'}

    def __init__(self, Q: Quiver, U: float=1, label: str="", angle: float=0,
                  coordinate: Literal['axes', 'figure', 'data', 'inches'] = 'axes',
                  labelpos: Literal['N', 'E', 'S', 'W'] = 'E',
                  labelsep: LenType=Pt(4),
                  loc: LocType = 'upper right',
                  bbox_to_anchor: tuple = None, color: any=None,
                  pad: LenType=Pt(5), borderpad: LenType=Pt(5),
                  frameon: bool=True, fontproperties=None,
                  style: Literal['default', 'legend'] = 'default',
                  **kwargs):
        """
        添加一个 Quiver 箭头图例，与 quiverkey 不同的是，像 legend 一样有边框且
        可以设置对齐方式。
        参数融合了 quiverkey 和 legend 的参数。

        Parameters
        ----------
        Q : Quiver
            `~.quiver()` 返回的对象。

            the return object of `~.quiver()`
        U : float, optional
            图例中箭头的长度，默认为 1。

            The length of the arrow in the legend, default is 1.
        label : str, optional
            图例中的文字说明。 The text in the legend.
        angle : float, optional
            箭头的角度，默认为 0。

            The angle of the arrow, default is 0.
        coordinate : Literal['axes', 'figure', 'data', 'inches'], optional
            图例所在的坐标系，默认为 'axes'。

            The coordinate system of the legend, default is 'axes'.
        labelpos : Literal['N', 'E', 'S', 'W'], optional
            文字在箭头的哪个方向，与 quiverkey 一致，默认为 'E'。

            The position of the text relative to the arrow, default is 'E'.
        labelsep : LenType, optional
            文字与箭头之间的间隔，单位为 pt，默认为 4 pt。

            The distance between the text and the arrow, default is 4 pt.
        loc : LocType, optional
            图例的位置，默认为 'upper right'。

            The location of the legend, default is 'upper right'.
        bbox_to_anchor : tuple, optional
            图例的锚点位置，默认为 None。

            The anchor point of the legend, default is None.
        color : any, optional
            箭头的颜色，默认跟随 quiver 的颜色。

            The color of the arrow, default is the same as the quiver.
        pad : LenType, optional
            图例的内边距，单位为字体大小，默认为 5 pt。

            The padding of the legend, default is 5 pt.
        borderpad : LenType, optional
            图例与锚点之间的间距，单位为字体大小，默认为 5 pt。

            The padding between the legend and the anchor point, default is 5 pt.
        frameon : bool, optional
            是否显示边框，默认为 True。

            Whether to show the border, default is True.
        fontproperties : any, optional
            字体属性，默认为 None。

            Font properties, default is None.
        style : Literal['default', 'legend'], optional
            预设样式，默认为白底无边框，可以设置为 'legend'，类似图例的样式。

            Preset style, default is 'default', can be set to 'legend' for a 
            legend-like style.

        **kwargs : any, optional
            其他设置边框的参数，传递给 `FancyBboxPatch`。

            Other parameters for the border, passed to `FancyBboxPatch`.

        """

        self.fig = Q.figure

        # 使用原来的 QuiverKey 类来创建箭头，为了保证适应不同的 MPL 版本
        self.qk = QuiverKey(Q, 0, 0, U, label, angle=angle, labelpos=labelpos,
                            coordinates=coordinate,)
        self.qk._init()
        qkwargs = Q.polykw
        self.vector = mcollections.PolyCollection(self.qk.verts, **qkwargs)
        self.vector.set_transform(Q.get_transform())
        self.vector.set_figure(self.fig)
        if color is not None:
            self.vector.set_color(color)

        # 设置文字
        self.text = mtext.Text( text=label, 
                               horizontalalignment=self.halign[labelpos],
                               verticalalignment=self.valign[labelpos],
                               fontproperties=fontproperties)
        self.text.set_figure(self.fig)

        labelsep = as_unit_value("pt", labelsep, fig=self.fig)
        text_offset = {"N": (0, labelsep),
                        "S": (0, -labelsep),
                        "E": (labelsep, 0),
                        "W": (-labelsep, 0)}[labelpos]
        self.text.set_position(np.array(text_offset)*self.dpi_cor)
        text_bbox = self._get_text_bbox()
        self.text.set_position(np.array(self.text.get_position())/self.dpi_cor)

        # 设置 Bbox
        vector_bbox = self._get_vector_bbox()
        legend_bbox = text_bbox.union([vector_bbox, text_bbox])

        self._da = DrawingArea(legend_bbox.width, legend_bbox.height,
                               max(-legend_bbox.x0, 0), max(-legend_bbox.y0, 0), clip=False)
        self.vector.set_offsets((0, 0))
        self.vector.set_offset_transform(self._da.get_transform())
        self._da.add_artist(self.vector)
        self._da.add_artist(self.text)

        zorder = kwargs.get('zorder', Q.get_zorder()+0.1)
        pad, borderpad = as_unit_value("fs", pad, borderpad, fig=self.fig)
        super().__init__(loc, child=self._da, pad=pad,
                         borderpad=borderpad, frameon=frameon,
                         bbox_to_anchor=bbox_to_anchor,
                         bbox_transform=self.qk.get_transform(), zorder=zorder)

        if style == 'legend':
            kwargs.setdefault('facecolor', self.fig.get_facecolor())
            kwargs.setdefault('edgecolor', '#666666')
            kwargs.setdefault('boxstyle', 'round, pad=0.3')
            kwargs.setdefault('alpha', rcParams['legend.framealpha'])
        elif style == 'default':
            kwargs.setdefault('facecolor', self.fig.get_facecolor())
            kwargs.setdefault('edgecolor', 'none')
        else:
            raise ValueError(f"Unknown style: {style}")
        self.patch.set(**kwargs)

    @property
    def dpi_cor(self):
        return self.fig.canvas.get_renderer().points_to_pixels(1.)

    def _get_vector_bbox(self) -> Bbox:
        vector_path = self.vector.get_paths()[0]
        xy_min = vector_path.vertices.min(axis=0)
        xy_max = vector_path.vertices.max(axis=0)
        xy_min = self.vector.get_transform().transform(xy_min)
        xy_max = self.vector.get_transform().transform(xy_max)
        return Bbox([xy_min/self.dpi_cor, xy_max/self.dpi_cor])

    def _get_text_bbox(self) -> Bbox:
        bbox = self.text.get_tightbbox(self.fig.canvas.get_renderer())
        return Bbox(bbox.get_points() / self.dpi_cor)


from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D
from scipy.interpolate import RegularGridInterpolator
from xarray import DataArray

class StreamQuiver(LineCollection):
    def __init__(self, X: np.ndarray | DataArray , Y: np.ndarray | DataArray,
                 U: np.ndarray | DataArray, V: np.ndarray | DataArray,
                 grid_num: int | tuple[int, int]=20, arrow_angle: float=30,
                 scale: float = 1, stagger: bool=True, head_length: float=0.3,
                 **kwargs):
        # 检查维度
        if isinstance(X, DataArray):
            X = X.values
        if isinstance(Y, DataArray):
            Y = Y.values
        if isinstance(U, DataArray):
            U = U.values
        if isinstance(V, DataArray):
            V = V.values

        if X.ndim != 1 or Y.ndim != 1:
            raise ValueError("X and Y must be 1D arrays.")
        if U.ndim != 2 or V.ndim != 2:
            raise ValueError("U and V must be 2D arrays.")

        self.X = X
        self.Y = Y
        self.U = U
        self.V = V

        if isinstance(grid_num, int):
            self.grid_num_x = grid_num
            self.grid_num_y = grid_num
        elif isinstance(grid_num, (tuple, list)) and len(grid_num) == 2:
            self.grid_num_x = grid_num[0]
            self.grid_num_y = grid_num[1]

        self.x_range = (X.min(), X.max())
        self.y_range = (Y.min(), Y.max())

        self.dx = (self.x_range[1] - self.x_range[0]) / (self.grid_num_x + 1)
        self.dy = (self.y_range[1] - self.y_range[0]) / (self.grid_num_y + 1)

        # 网格中心点位置
        grid_x = np.linspace(self.x_range[0], self.x_range[1], self.grid_num_x)
        grid_y = np.linspace(self.y_range[0], self.y_range[1], self.grid_num_y)
        self.grid_x, self.grid_y = np.meshgrid(grid_x, grid_y, indexing='ij')
        if stagger:
            # 奇数行错开
            self.grid_x += 0.1547 * (self.grid_x - self.x_range[0])
            self.grid_x[:, 1::2] += self.dx / 2
            valid = self.grid_x < self.x_range[1]
            self.grid_x = self.grid_x[valid]
            self.grid_y = self.grid_y[valid]

        self.kwargs = kwargs
        self.arrow_angle = np.deg2rad(arrow_angle)
        self.steps = 5
        self.dt = scale
        self.head_length = head_length
        self.lines = self.generate_lines()


    def rk2(self, x, y, interp_u, interp_v, inverse=False):
        """RK2 流线积分"""
        dt = -self.dt if inverse else self.dt
        k1u = dt * interp_u((y, x))
        k1v = dt * interp_v((y, x))
        k2u = dt * interp_u((y + k1v, x + k1u))
        k2v = dt * interp_v((y + k1v, x + k1u))
        x_new = x + 0.5 * (k1u + k2u)
        y_new = y + 0.5 * (k1v + k2v)
        return x_new, y_new


    def generate_lines(self):
        #  每个格点都积分流线
        interp_u = RegularGridInterpolator((self.Y, self.X), self.U,
                                           bounds_error=False, fill_value=None)
        interp_v = RegularGridInterpolator((self.Y, self.X), self.V,
                                           bounds_error=False, fill_value=None)
        x = self.grid_x.flatten()
        y = self.grid_y.flatten()
        line_x = np.zeros((len(x), self.steps*2+1))
        line_y = np.zeros((len(y), self.steps*2+1))
        line_x[:, self.steps] = x
        line_y[:, self.steps] = y
        line_length = np.zeros((len(x)))

        grid_u = interp_u((y, x))
        grid_v = interp_v((y, x))
        grid_speed = np.sqrt(grid_u**2 + grid_v**2)

        # 向前积分
        for k in range(self.steps):
            x_new, y_new = self.rk2(x, y, interp_u, interp_v, inverse=False)
            line_x[:, self.steps + k + 1] = x_new
            line_y[:, self.steps + k + 1] = y_new
            line_length += np.sqrt((x_new - x)**2 + (y_new - y)**2)
            x, y = x_new, y_new

        # 向后积分
        x = self.grid_x.flatten()
        y = self.grid_y.flatten()
        for k in range(self.steps):
            x_new, y_new = self.rk2(x, y, interp_u, interp_v, inverse=True)
            line_x[:, self.steps - k - 1] = x_new
            line_y[:, self.steps - k - 1] = y_new
            line_length += np.sqrt((x_new - x)**2 + (y_new - y)**2)
            x, y = x_new, y_new

        # 画线
        lines = []
        for i in range(len(line_x)):
            x = line_x[i]
            y = line_y[i]

            # 画箭头
            ang = np.arctan2(y[-1]-y[-2], x[-1]-x[-2])
            head_x = [0, x[-1], 0]
            head_y = [0, y[-1], 0]
            head_length = self.dx * self.head_length
            if line_length[i] < head_length:
                head_length = line_length[i]
            head_x[0] = x[-1] - head_length * np.cos(ang - self.arrow_angle)
            head_y[0] = y[-1] - head_length * np.sin(ang - self.arrow_angle)
            head_x[2] = x[-1] - head_length * np.cos(ang + self.arrow_angle)
            head_y[2] = y[-1] - head_length * np.sin(ang + self.arrow_angle)

            lines.append(np.array([line_x[i], line_y[i]]).T)
            lines.append(np.array([head_x, head_y]).T)

        super().__init__(lines, **self.kwargs)

    def quiverkey(self, ax: Axes, X=0.8, Y=1.05, U: float=1, 
                  label: str="",
                  labelpos: Literal['N', 'E', 'S', 'W'] = 'E',
                  pad:LenType=Pt(5),
                  **kwargs):
        arrow_length = self.dt * U * (self.steps * 2)
        arrow = np.array([[0, 0], [arrow_length, 0]])
        # 画箭头
        ang = 0
        head_x = [0, arrow_length, 0]
        head_y = [0, 0, 0]
        head_length = self.dx * self.head_length
        if arrow_length < head_length:
            head_length = arrow_length
        head_x[0] = arrow_length - head_length * np.cos(ang - self.arrow_angle)
        head_y[0] = 0 - head_length * np.sin(ang - self.arrow_angle)
        head_x[2] = arrow_length - head_length * np.cos(ang + self.arrow_angle)
        head_y[2] = 0 - head_length * np.sin(ang + self.arrow_angle)
        arrow_head = np.array([head_x, head_y]).T

        # 转为 Axes 坐标
        arrow = ax.transData.transform(arrow)
        arrow = ax.transAxes.inverted().transform(arrow)
        arrow_head = ax.transData.transform(arrow_head)
        arrow_head = ax.transAxes.inverted().transform(arrow_head)
        origin = arrow[0]
        arrow_head -= origin
        arrow -= origin
        arrow[:, 0] += X
        arrow[:, 1] += Y
        arrow_head[:, 0] += X
        arrow_head[:, 1] += Y

        args = self.kwargs.copy()
        args.update(transform=ax.transAxes, clip_on=False)
        line1 = Line2D(arrow[:, 0], arrow[:, 1], **args)
        line2 = Line2D(arrow_head[:, 0], arrow_head[:, 1], **args)
        ax.add_artist(line1)
        ax.add_artist(line2)

        # 文字
        halign = {'N': 'center', 'S': 'center', 'E': 'left', 'W': 'right'}
        valign = {'N': 'bottom', 'S': 'top', 'E': 'center', 'W': 'center'}
        if labelpos in ['E', 'W']:
            pad = as_unit_value("relx", pad, fig=ax.figure)
        else:
            pad = as_unit_value("rely", pad, fig=ax.figure)
        arrow_center = (arrow[0] + arrow[-1]) / 2
        text_anchor = {
            "N": (arrow_center[0], arrow_center[1] + pad),
            "S": (arrow_center[0], arrow_center[1] - pad),
            "E": (arrow[1, 0] + pad, arrow[1, 1]),
            "W": (arrow[0, 0] - pad, arrow[1, 1]),
        }
        ax.text(x=text_anchor[labelpos][0],
                          y=text_anchor[labelpos][1],
                          s=label,
                          horizontalalignment=halign[labelpos],
                          verticalalignment=valign[labelpos],
                          transform=ax.transAxes, clip_on=False,
                          **kwargs)
        return text_anchor[labelpos]


class StreamQuiver2(LineCollection):
    def __init__(self,fig, ax: Axes, X: np.ndarray | DataArray , Y: np.ndarray | DataArray,
                 U: np.ndarray | DataArray, V: np.ndarray | DataArray,
                 numbers: int | tuple[int, int] = 15,
                 pivot: Literal['tail', 'middle', 'tip'] = 'middle',
                 scale: float = 1, stagger: bool=True, steps: int=4,
                 head_angle: float=30, head_length: float= 1, **kwargs
                 ):

        self.X = X
        self.Y = Y
        self.U = U
        self.V = V
        self.ax = ax

        # only 1D arrays for X and Y can use RegularGridInterpolator
        if X.ndim != 1 or Y.ndim != 1:
            raise ValueError("X and Y must be 1D arrays.")

        if U.ndim != 2 or V.ndim != 2:
            raise ValueError("U and V must be 2D arrays.")

        if isinstance(numbers, int):
            self.nx = numbers
            self.ny = numbers
        elif isinstance(numbers, tuple) and len(numbers) == 2:
            self.nx = numbers[0]
            self.ny = numbers[1]
        else:
            raise ValueError("numbers must be int or tuple of two ints.")

        self.pivot = pivot
        self.steps = steps
        self.kwargs = kwargs
        self.head_angle = np.deg2rad(head_angle)
        self.head_length = head_length

        # data to axes transform
        self.tranAxes = ax.transData + ax.transAxes.inverted()
        self.tranAxesInv = self.tranAxes.inverted()
        
        # generate grid points
        self.x_range = (X.min(), X.max())
        self.y_range = (Y.min(), Y.max())

        self.dx = (self.x_range[1] - self.x_range[0]) / (self.nx - 1)
        self.dy = (self.y_range[1] - self.y_range[0]) / (self.ny - 1)

        # 网格中心点位置
        grid_x = np.linspace(self.x_range[0], self.x_range[1], self.nx)
        grid_y = np.linspace(self.y_range[0], self.y_range[1], self.ny)

        self.grid_x, self.grid_y = np.meshgrid(grid_x, grid_y, indexing='ij')

        if stagger:
            self.grid_x[:, 1::2] += self.dx / 2
            valid = self.grid_x < self.x_range[1]
            self.grid_x = self.grid_x[valid]
            self.grid_y = self.grid_y[valid]
        
        self.dt = scale

        self.draw_lines()
        

    def rk2(self, x, y, interp_u, interp_v, inverse=False):
        """RK2 流线积分"""
        dt = -self.dt if inverse else self.dt
        k1u = dt * interp_u((y, x))
        k1v = dt * interp_v((y, x))
        k2u = dt * interp_u((y + k1v, x + k1u))
        k2v = dt * interp_v((y + k1v, x + k1u))
        x_new = x + 0.5 * (k1u + k2u)
        y_new = y + 0.5 * (k1v + k2v)
        return x_new, y_new
    

    def draw_lines(self):
        #  每个格点都积分流线
        interp_u = RegularGridInterpolator((self.Y, self.X), self.U,
                                           bounds_error=False, fill_value=None)
        interp_v = RegularGridInterpolator((self.Y, self.X), self.V,
                                           bounds_error=False, fill_value=None)
        x = self.grid_x.flatten()
        y = self.grid_y.flatten()
        line_x = np.zeros((len(x), self.steps*2+1))
        line_y = np.zeros((len(y), self.steps*2+1))

        pivot_idx = {'tail': 0, 'middle': self.steps, 'tip': self.steps*2}[self.pivot]
        line_x[:, pivot_idx] = x
        line_y[:, pivot_idx] = y
        line_length = np.zeros((len(x)))

        # grid_u = interp_u((y, x))
        # grid_v = interp_v((y, x))
        # grid_speed = np.sqrt(grid_u**2 + grid_v**2)
        
        forward_steps = {'tail': self.steps*2, 'middle': self.steps, 'tip': 0}[self.pivot]
        backward_steps = self.steps*2 - forward_steps

        # 向前积分
        for k in range(forward_steps):
            x_new, y_new = self.rk2(x, y, interp_u, interp_v, inverse=False)
            line_x[:, pivot_idx + k + 1] = x_new
            line_y[:, pivot_idx + k + 1] = y_new
            line_length += np.sqrt((x_new - x)**2 + (y_new - y)**2)
            x, y = x_new, y_new

        # 向后积分
        x = self.grid_x.flatten()
        y = self.grid_y.flatten()
        for k in range(backward_steps):
            x_new, y_new = self.rk2(x, y, interp_u, interp_v, inverse=True)
            line_x[:, pivot_idx - k - 1] = x_new
            line_y[:, pivot_idx - k - 1] = y_new
            line_length += np.sqrt((x_new - x)**2 + (y_new - y)**2)
            x, y = x_new, y_new

        # 画线
        lines = []
        for i in range(len(line_x)):
            x = line_x[i, -2:]
            y = line_y[i, -2:]
            x, y = self.tranAxes.transform(np.array([x, y]).T).T

            # 画箭头
            ang = np.arctan2(y[-1]-y[-2], x[-1]-x[-2])
            head_x = [0, x[-1], 0]
            head_y = [0, y[-1], 0]
            head_length = self.head_length
            # if line_length[i] < head_length:
            #     head_length = line_length[i] * 0.8
            head_x[0] = x[-1] - head_length * np.cos(ang - self.head_angle)
            head_y[0] = y[-1] - head_length * np.sin(ang - self.head_angle)
            head_x[2] = x[-1] - head_length * np.cos(ang + self.head_angle)
            head_y[2] = y[-1] - head_length * np.sin(ang + self.head_angle)

            head_x, head_y = self.tranAxesInv.transform(np.array([head_x, head_y]).T).T

            lines.append(np.array([line_x[i], line_y[i]]).T)
            lines.append(np.array([head_x, head_y]).T)

        super().__init__(lines, **self.kwargs)